;;;----  Gauss elimination method to Solve Linear Eqns
;;Problem: solve A x = b, where A is invertible
;;Steps:
;;1. By row transformations to rescale each row's maximum into be +/- 1;
;;2. Gauss elimination
;;3. Get the solution

;; load:
(load "matrix-vector-operations.scm")


;;; for step 1. A,b => A',b' (就地修改)
(define (rescale-A-and-b! A b n)
  (define (abs-max v) (apply max (map abs (vector->list v))))
  (for-each (lambda (i)
              (let [(absmax (abs-max (vector-ref A i)))]
                (vector-map! (lambda (ele) (/ ele absmax)) (vector-ref A i))
                (vector-set! b i (/ (vector-ref b i) absmax))))
            (range n)))


;;;for step 2. Gauss elimination

;; in A_jk, j>=k; to get main element(absmax) and it's place
(define (find-main-element A k n);;=> the place of the main-element
  (let [(index k)
        (absmax (abs (matrix-ref A k k)))]
    (for-each (lambda (j)
                (if (> (abs (matrix-ref A j k)) absmax)
                    (begin (set! index j)
                           (set! absmax (matrix-ref A j k)))))
              (range (+ k 1) n))
    index))

;; swap! ith line and jth line elements whose index >=k
(define (matrix-swap! A i j k n)
  (for-each (lambda (l)
              (let [(t 0)]
                (set! t (matrix-ref A i l))
                (matrix-set! A i l (matrix-ref A j l))
                (matrix-set! A j l t)))
            (range k n)))


(define (Gauss-elimination! A b)
  (define n (vector-length A))
  ;;1. rescale! A and b
  (rescale-A-and-b! A b n)
  ;;2. elimination! A and b
  (for-each (lambda (k)
              ;; swap! the main-element line to top
              (let [(ik (find-main-element A k n))]
                (if (not (= ik k)) (begin (matrix-swap! A k ik k n)
                                          (vector-swap! b k ik))))
              ;; get A_kk
              (define akk (matrix-ref A k k))
              (if (= akk 0) (error "The matrix is not invertible!"))
              ;; change! A_ij with j>=k,i>=k+1; and b_i with i>=k+1
              (for-each (lambda (i)
                          ;; A_ij with j=k,i>=k+1: A_ik <- A_ik / A_kk
                          (matrix-set! A i k (/ (matrix-ref A i k) akk))
                          ;; A_ij with j,i>=k+1: A_ij <- A_ij - A_ik * A_kj
                          (for-each (lambda (j)
                                      (matrix-set! A i j
                                                   (- (matrix-ref A i j)
                                                      (* (matrix-ref A i k) (matrix-ref A k j)))))
                                    (range (+ k 1) n))
                          ;; b_j <- b_j - A_ik * b_k
                          (vector-set! b i (- (vector-ref b i)
                                              (* (matrix-ref A i k) (vector-ref b k)))))
                        (range (+ k 1) n)))
            (range (- n 1)))
  ;;3. get the result
  (define result (make-vector n 0))
  (if (= 0 (matrix-ref A (- n 1) (- n 1))) (error "The matrix is not invertible!"))
  (vector-set! result (- n 1) (/ (vector-ref b (- n 1)) (matrix-ref A (- n 1) (- n 1))))
  (for-each (lambda (k)
              (vector-set! result k
                           (/ (- (vector-ref b k)
                                 (let [(sum 0)]
                                   (for-each (lambda (j)
                                               (set! sum (+ sum
                                                            (* (matrix-ref A k j)
                                                               (vector-ref result j)))))
                                             (range (+ k 1) n))
                                   sum))
                              (matrix-ref A k k))))
            (range (- n 2) -1))
  result)

;;;------------ test
(let [(A (vector #(2 2 3) #(4 7 7) #(-2 4 5)))
      (b #(3 1 -7))]
  (newline)
  (print "precise solution: " #(2 -2 1))
  (print "solution here:    " (Gauss-elimination! (matrix-copy A)
                                                  (vector-copy b))))


